import torch
from torch import nn
import torch.nn.functional as F


# 模型
class Model(nn.Module):
    def __init__(self, embedding_pretrained, num_filters, filter_sizes, dropout, num_classes):
        super(Model, self).__init__()
        embed = embedding_pretrained.size()[1]
        self.embedding = nn.Embedding.from_pretrained(embedding_pretrained, freeze=False)
        self.convs = nn.ModuleList(
            [nn.Conv2d(1, num_filters, (k, embed)) for k in filter_sizes])
        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(num_filters * len(filter_sizes), num_classes)

    def conv_and_pool(self, x, conv):
        x = F.relu(conv(x)).squeeze(3)
        x = F.max_pool1d(x, x.size(2)).squeeze(2)
        return x

    def forward(self, x):
        out = self.embedding(x)
        out = out.unsqueeze(1)

        out = torch.cat([self.conv_and_pool(out, conv) for conv in self.convs], 1)
        out = self.dropout(out)
        out = self.fc(out)
        return out
